package base

import (
	"context"
	"strings"

	"gitee.com/dn-jinmin/chatcms/internal/svc"
	"gitee.com/dn-jinmin/langchaingo"
	"gitee.com/dn-jinmin/langchaingo/chainx"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/prompts"
)

type Base struct {
	svc    *svc.ServiceContext
	llm    llms.Model
	chains *chainx.LLMChain
}

func NewBase(svc *svc.ServiceContext, template prompts.PromptTemplate) *Base {
	return &Base{
		svc:    svc,
		llm:    svc.LLm,
		chains: chainx.NewLLMChain(svc.LLm, template, chains.WithCallback(svc.Callback)),
	}
}

func (b *Base) ChainOption(ctx context.Context, input map[string]any) []chainx.LLmChainsOption {
	//model := input[InputModel].(string)
	model := ctx.Value(InputModel).(string)
	return []chainx.LLmChainsOption{
		chainx.LLmChainsModel(model),
	}
}

// 确定模型是否思考
func (b *Base) Think(ctx context.Context, input map[string]any) map[string]any {
	// 选择模型
	model, ok := ctx.Value(InputModel).(string)
	if !ok {
		// 默认通义千问3
		return input
	}

	// 通义千问3 模型需要特殊处理
	if !strings.Contains(model, Qwen) {
		input[InputThink] = ""
		return input
	}

	// 判断是否需要 /no_think
	t, ok := ctx.Value(InputThink).(bool)
	if !ok {
		return input
	}
	if !t {
		input[InputThink] = ""
	}

	return input
}

func (b *Base) Call(ctx context.Context, input map[string]any) (map[string]any, error) {
	// 调用模型
	return b.chains.Call(ctx, b.Think(ctx, input), b.ChainOption(ctx, input)...)
}

func NewPrompt(tmp string, inputs ...string) prompts.PromptTemplate {
	return prompts.NewPromptTemplate(tmp, append([]string{
		langchaingo.Input,
		InputThink,
	}, inputs...))
}
